/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.sdk.extensions.sql.impl.rel;



import com.bff.gaia.unified.sdk.Pipeline;

import com.bff.gaia.unified.sdk.transforms.PTransform;

import com.bff.gaia.unified.sdk.values.PCollection;

import com.bff.gaia.unified.sdk.values.PCollectionList;

import com.bff.gaia.unified.sdk.values.Row;

import org.apache.calcite.plan.volcano.RelSubset;

import org.apache.calcite.rel.RelNode;



import java.util.HashMap;

import java.util.List;

import java.util.Map;

import java.util.stream.Collectors;



/** Utilities for {@code UnifiedRelNode}. */

public class UnifiedSqlRelUtils {



  public static PCollection<Row> toPCollection(Pipeline pipeline, UnifiedRelNode node) {

    return toPCollection(pipeline, node, new HashMap());

  }



  /** Transforms the inputs into a PInput. */

  private static PCollectionList<Row> buildPCollectionList(

	  List<RelNode> inputRels, Pipeline pipeline, Map<Integer, PCollection<Row>> cache) {

    if (inputRels.isEmpty()) {

      return PCollectionList.empty(pipeline);

    } else {

      return PCollectionList.of(

          inputRels.stream()

              .map(input -> UnifiedSqlRelUtils.toPCollection(pipeline, (UnifiedRelNode) input, cache))

              .collect(Collectors.toList()));

    }

  }



  /**

   * A {@link UnifiedRelNode} is a recursive structure, the {@code UnifiedQueryPlanner} visits it with a

   * DFS(Depth-First-Search) algorithm.

   */

  static PCollection<Row> toPCollection(

	  Pipeline pipeline, UnifiedRelNode node, Map<Integer, PCollection<Row>> cache) {

    PCollection<Row> output = cache.get(node.getId());

    if (output != null) {

      return output;

    }



    String name = node.getClass().getSimpleName() + "_" + node.getId();

    PCollectionList<Row> input = buildPCollectionList(node.getPCollectionInputs(), pipeline, cache);

    PTransform<PCollectionList<Row>, PCollection<Row>> transform = node.buildPTransform();

    output = Pipeline.applyTransform(name, input, transform);



    cache.put(node.getId(), output);

    return output;

  }



  public static UnifiedRelNode getUnifiedRelInput(RelNode input) {

    if (input instanceof RelSubset) {

      // go with known best input

      input = ((RelSubset) input).getBest();

    }

    return (UnifiedRelNode) input;

  }

}